#!/bin/sh

set -e

# The target version
# Make sure you have download the package
#   in the same path of this script
ES_UPGRADE_TO=${ES_UPGRADE_TO:-7.17.25}

# The URL of ES
ES_URL=https://192.168.56.11:9200

# The superuser name of ES
ES_SUPER_NAME=${ES_SUPER_NAME:-elastic}

# Log path
ES_LOG_FILE=/var/log/elasticsearch/uat.log

# Get script dir
DIRNAME=`realpath $0 | xargs dirname`
CURL_PREFIX="curl -sk -u${ES_SUPER_NAME}:${ES_SUPER_PASS}"

function help () {
    cat es_upgrade.readme
}

function pre_check () {
    if [[ -z "${ES_SUPER_PASS}" ]]
    then
        echo "
            ES_SUPER_PASS is not defined
            Please set 'export ES_SUPER_PASS=<password>' first
            Example:
                set +o history
                export ES_SUPER_PASS=A0DAndkasfd0
                set -o history
        "
        exit
    fi

    ES_STATUS=`pgrep -af java | grep elasticsearch || echo "Stopped"`
    if [[ "${ES_STATUS}" != "Stopped" ]]
    then
        ${CURL_PREFIX} ${ES_URL}/_cluster/health?pretty
    fi
}

function pre_tasks () {
    echo -e "\nStop non-essential indexing and perform a flush"
    for i in {1..99}
    do
        set -e
        FLUSH_INFO=`${CURL_PREFIX} -XPOST ${ES_URL}/_flush?pretty`
        TOTAL=`echo ${FLUSH_INFO} | awk '{print $7}' | sed 's/,//'`
        SUCCESSFUL=`echo ${FLUSH_INFO} | awk '{print $10}' | sed 's/,//'`
        FAILED=`echo ${FLUSH_INFO} | awk '{print $13}' | sed 's/,//'`

        if [[ "$FAILED" -eq 0 ]] && [[ "$SUCCESSFUL" -eq "$TOTAL" ]]
        then
            echo '  Flushed!'
            break
        else
            echo '  Flushing...'
            sleep 15
        fi
    done

    echo -e "\nTemporarily stop the tasks associated with active machine learning jobs and datafeeds"
    ${CURL_PREFIX} -XPOST "${ES_URL}/_ml/set_upgrade_mode?enabled=true&pretty" | grep 'true'

    echo -e "\nDisable shard allocation"
    ${CURL_PREFIX} -XPUT ${ES_URL}/_cluster/settings?pretty -H 'Content-Type: application/json' -d'
    {
        "persistent": {
            "cluster.routing.allocation.enable": "primaries"
        }
    }' | grep 'true'
}

function post_tasks () {
    echo -e "\nReenable shard allocation"
    ${CURL_PREFIX} -XPUT ${ES_URL}/_cluster/settings?pretty -H 'Content-Type: application/json' -d'
    {
        "persistent": {
            "cluster.routing.allocation.enable": null
        }
    }'

    echo -e "\nResume the tasks associated with active machine learning jobs and datafeeds"
    ${CURL_PREFIX} -XPOST "${ES_URL}/_ml/set_upgrade_mode?enabled=false&pretty" | grep 'true'
}

function es_stop () {
    pre_check

    pre_tasks

    service elasticsearch stop

    for i in {1..99}
    do
        ES_STATUS=`pgrep -af java | grep elasticsearch || echo "Stopped"`

        if [[ "${ES_STATUS}" -eq "Stopped" ]]
        then
            echo "Service is stopped"
            break
        else
            echo "Service is stopping..."
            sleep 15
        fi
    done
}

function es_start () {
    pre_check

    ES_STATUS=`pgrep -af java | grep elasticsearch || echo "Stopped"`

    service elasticsearch start

    if [[ "${ES_STATUS}" == "Stopped" ]]
    then
        tail -F -n50 ${ES_LOG_FILE} | sed '/9200/q' || /bin/true
    fi

    post_tasks

    echo -e "\nWait for cluster turn into healthy"
    for i in {1..99}
    do
        ES_HEALTH=`${CURL_PREFIX} "${ES_URL}/_cluster/health?pretty" | egrep '(green|yellow)' || echo "Not ready"`

        if [[ "${ES_HEALTH}" == "Not ready" ]]
        then
            echo ${ES_HEALTH}
            sleep 15
        else
            echo "Cluster is healthy now"
            break
        fi
    done
}

function es_upgrade () {
    rpm -Uvh ${DIRNAME}/elasticsearch-${ES_UPGRADE_TO}-x86_64.rpm || /bin/true

    systemctl daemon-reload

    systemctl enable elasticsearch

    /usr/share/elasticsearch/bin/elasticsearch --version
}

case "$@" in
    stop)
        es_stop
        ;;
    start)
        es_start
        ;;
    upgrade)
        es_stop
        es_upgrade
        es_start
        ;;
    restart)
        es_stop
        sleep 10
        es_start
        ;;
    help | -h | --help | *)
        help
        ;;
esac
